/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package asset.load;

import com.jme3.math.Vector2f;
import com.jme3.math.Vector3f;
import com.jme3.scene.Mesh;
import com.jme3.scene.VertexBuffer;
import com.jme3.util.BufferUtils;
import asset.MeshSet;
import asset.Skel;
import java.io.File;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import org.w3c.dom.Document;
import org.w3c.dom.Node;
import util.math.Convert;
import util.xml.XML;

/**
 *
 * @author
 */
public class LoadMeshSet {

    public static MeshSet load(Skel skel, String path) throws Exception {
        return load(skel, new File(path));
    }
    
    /**
     * 
     * @param skel provided only to translate bone names to bone indices
     * @param file
     * @return
     * @throws Exception 
     */
    public static MeshSet load(Skel skel, File file) throws Exception {
        Worker w = new Worker(skel, file);
        w.load();
        return w.getMeshSet();
    }

    private static class Worker {

        Skel skel;
        File file;
        //
        MeshSet meshSet;

        private Worker(Skel skel, File file) {
            this.skel = skel;
            this.file = file;
        }

        public MeshSet getMeshSet() {
            return meshSet;
        }

        public void load() throws Exception {
            Document doc = XML.openDocument(file);

            //---------------------------------------------
            // Parse vertex positions, normals, and texture coords

            Node verticesNode = doc.getFirstChild().getFirstChild();

            int numVerts = XML.getInt(verticesNode, "count");

            List<VertexBuffer> sharedBuffers = new ArrayList<>();

            parseVertices(verticesNode, numVerts, sharedBuffers);

            //---------------------------------------------
            // Parse submeshes: we map the name of each submesh to the corresponding 
            // indices buffer.
            Map<String, VertexBuffer> submeshes = new HashMap<>();

            Node submeshesNode = verticesNode.getNextSibling();

            parseSubmeshes(submeshesNode, submeshes);

            //---------------------------------------------
            // Parse bone weights.   
            Node skinsNode = submeshesNode.getNextSibling();
            parseSkin(skinsNode, numVerts, sharedBuffers);

            //---------------------------------------------
            // Create the geometries
            this.meshSet = new MeshSet();
            for (Map.Entry<String, VertexBuffer> e : submeshes.entrySet()) {
                String name = e.getKey();
                VertexBuffer indices = e.getValue();

                Mesh mesh = new Mesh();
                // Add the shared buffers
                for (VertexBuffer b : sharedBuffers) {
                    mesh.setBuffer(b);
                }

                // Add the indices buffer
                mesh.setBuffer(indices);
                
                mesh.updateBound();
                
                // Some settings for animation
                mesh.setMaxNumWeights(4);
                mesh.generateBindPose(true);
                
                this.meshSet.addSubmesh(name, mesh);
            }
        }

        private void parseVertices(Node verticesNode, int numVerts, List<VertexBuffer> sharedBuffers) {
            Vector3f[] positions = new Vector3f[numVerts];
            Vector3f[] normals = new Vector3f[numVerts];
            Vector2f[] texCoord = new Vector2f[numVerts];

            // Get all vertex data
            int i = 0;
            for (Node n : XML.iterate(verticesNode.getChildNodes())) {
                Node coordsNode = n.getFirstChild();
                positions[i] = Convert.toVector3f(XML.getFloatArray(coordsNode, "x", "y", "z"));

                Node normalNode = coordsNode.getNextSibling();
                normals[i] = Convert.toVector3f(XML.getFloatArray(normalNode, "x", "y", "z"));

                Node texNode = normalNode.getNextSibling();
                texCoord[i] = Convert.toVector2f(XML.getFloatArray(texNode, "u", "v"));

                ++i;
            }

            VertexBuffer bufPositions = new VertexBuffer(VertexBuffer.Type.Position);
            bufPositions.setupData(VertexBuffer.Usage.Static, 3, VertexBuffer.Format.Float, BufferUtils.createFloatBuffer(positions));
            sharedBuffers.add(bufPositions);

            VertexBuffer bufNormals = new VertexBuffer(VertexBuffer.Type.Normal);
            bufNormals.setupData(VertexBuffer.Usage.Static, 3, VertexBuffer.Format.Float, BufferUtils.createFloatBuffer(normals));
            sharedBuffers.add(bufNormals);

            VertexBuffer bufTexCoord = new VertexBuffer(VertexBuffer.Type.TexCoord);
            bufTexCoord.setupData(VertexBuffer.Usage.Static, 2, VertexBuffer.Format.Float, BufferUtils.createFloatBuffer(texCoord));
            sharedBuffers.add(bufTexCoord);

        }

        private void parseSubmeshes(Node submeshesNode, Map<String, VertexBuffer> submeshes) {
            for (Node n : XML.iterate(submeshesNode.getChildNodes())) {
                List<Integer> indices = new ArrayList<>();

                // Go over the polygons
                for (Node p : XML.iterate(n.getChildNodes())) {
                    // Go over the vertices in this polygon.
                    // We make no assumption about how many vertices a polygon contains.
                    // We convert the polygon to faces on the fly. The user is responsible
                    // for providing "convex" polygons, so our generated faces look
                    // reasonable.
                    // We assume the vertices are given counter-clockwise in the data.

                    Iterator<Node> iter = XML.iterate(p.getChildNodes()).iterator();

                    // All faces are generated from the first vertex.
                    int fixedVertexId = XML.getInt(iter.next(), "id");
                    int prev = XML.getInt(iter.next(), "id");
                    while (iter.hasNext()) {
                        int cur = XML.getInt(iter.next(), "id");

                        indices.add(fixedVertexId);
                        indices.add(prev);
                        indices.add(cur);

                        prev = cur;
                    }
                }

                // Put indices in array
                int[] _indices = new int[indices.size()];
                int j = 0;
                for (int v : indices) {
                    _indices[j++] = v;
                }

                VertexBuffer bufIndices = new VertexBuffer(VertexBuffer.Type.Index);
                bufIndices.setupData(VertexBuffer.Usage.Static, 3, VertexBuffer.Format.UnsignedInt, BufferUtils.createIntBuffer(_indices));

                String name = XML.getStr(n, "name");
                submeshes.put(name, bufIndices);
            }
        }

        //======================================================================
        class BoneWeight implements Comparable<BoneWeight> {

            public int boneIndex;
            public float weight;

            public BoneWeight(int boneIndex, float weight) {
                this.boneIndex = boneIndex;
                this.weight = weight;
            }

            @Override
            public int compareTo(BoneWeight o) {
                if (this.weight > o.weight) {
                    return -1;
                } else if (this.weight < o.weight) {
                    return 1;
                } else {
                    return 0;
                }
            }
        }

        class W {

            List<BoneWeight> weights;

            public W() {
                weights = new ArrayList<>();
            }

            public void add(int boneIndex, float weight) {
                weights.add(new BoneWeight(boneIndex, weight));
            }
        }

        private void parseSkin(Node skinsNode, int numVerts, List<VertexBuffer> sharedBuffers) {
            
            // We cannot generate bone indices when there is no skeleton
            // given.
            if(skel == null) {
                return;
            }

            // First we collect the weights for each vertex.         
            W[] vertWeights = new W[numVerts];
            for (int j = 0; j < numVerts; ++j) {
                vertWeights[j] = new W();
            }


            if (skinsNode.getChildNodes().getLength() > 0) {
                Node skinNode = skinsNode.getFirstChild(); // we only process the first weight map

                for (Node boneNode : XML.iterate(skinNode.getChildNodes())) {
                    String boneName = XML.getStr(boneNode, "name");
                    int boneIndex = skel.getBoneIndex(boneName);

                    for (Node vertNode : XML.iterate(boneNode.getChildNodes())) {
                        int id = XML.getInt(vertNode, "id");
                        float weight = XML.getFloat(vertNode, "weight");

                        vertWeights[id].add(boneIndex, weight);
                    }
                }
            }

            float[] boneWeights = new float[numVerts * 4];
            byte[] boneIndices = new byte[numVerts * 4]; // we use bytes (like in the ogre mesh loader)

            // Now we select for each vertex the four strongest bones, and we
            // reweight them (normalization).            

            for (int j = 0; j < numVerts; ++j) {
                W weights = vertWeights[j];
                Collections.sort(weights.weights);

                // Compute sum
                float weightSum = 0.0f;
                Iterator<BoneWeight> iter = weights.weights.iterator();
                int count = 0;
                while (count < 4 && iter.hasNext()) {
                    BoneWeight bw = iter.next();
                    weightSum += bw.weight;
                    count += 1;
                }

                // Iterate again
                iter = weights.weights.iterator();
                int k = 0;
                while (k < count && iter.hasNext()) {
                    BoneWeight bw = iter.next();

                    float w = bw.weight / weightSum;

                    boneWeights[j * 4 + k] = w;
                    boneIndices[j * 4 + k] = (byte) bw.boneIndex;
                    k += 1;
                }

                // If the vertex does not have 4 bones, we give the remaining
                // slots zero weight
                while (k < 4) {
                    boneWeights[j * 4 + k] = 0.0f;
                    boneIndices[j * 4 + k] = (byte) 0; // just an arbitrary bone index (the weight is zero anyway)
                    k += 1;
                }
            }

            VertexBuffer bufBoneWeights = new VertexBuffer(VertexBuffer.Type.BoneWeight);
            bufBoneWeights.setupData(VertexBuffer.Usage.Static, 4, VertexBuffer.Format.Float, BufferUtils.createFloatBuffer(boneWeights));
            sharedBuffers.add(bufBoneWeights);

            VertexBuffer bufBoneIndices = new VertexBuffer(VertexBuffer.Type.BoneIndex);
            bufBoneIndices.setupData(VertexBuffer.Usage.Static, 4, VertexBuffer.Format.UnsignedByte, BufferUtils.createByteBuffer(boneIndices));
            sharedBuffers.add(bufBoneIndices);
        }
    }
}
